function [P,Q,W,X,Y,Z] = BPTF2(train_data, test_data, D, max_epoch, init_feat, tag)
% train_data - training data
% test_data - testing data
% D - num of features
% max_epoch - number of samples 
% init_feat - MAP estimates of feature vectors 

alpha = 2;

N = train_data.N;
M = train_data.M;
K = train_data.K;
mean_rating = train_data.mean_rating; 

filterU = train_data.filterU; 
filterI = train_data.filterI; 
filterT = train_data.filterT; 

test = test_data.probe;

if isempty(init_feat)
	fprintf('Initializing feature vectors randomly\n');
	P = 0.1*randn(D,N);
	Q = 0.1*randn(D,M);

	W = 0.1*randn(D,M);
	X = 0.1*randn(D,K);

	Y = 0.1*randn(D,K);
	Z = 0.1*randn(D,N);
else
	fprintf('Initializing feature vectors with MAP estimates\n');
	P = init_feat.P;
	Q = init_feat.Q;

	W = init_feat.W;
	X = init_feat.X;
	
	Y = init_feat.Y;
	Z = init_feat.Z;
end

err = zeros(max_epoch, 1);
all_ratings = predict(test, P, Q, W, X, Y, Z);
all_ratings = all_ratings + mean_rating;
count = 1;
tic;

for epoch=1:max_epoch
    [mu_P, L_P] = rnd_user_hyperparams(P);
    [mu_Q, L_Q] = rnd_item_hyperparams(Q);
    
    [mu_W, L_W] = rnd_item_hyperparams(W);
    [mu_X, L_X] = rnd_time_hyperparams(X);
    
    [mu_Y, L_Y] = rnd_time_hyperparams(Y);
    [mu_Z, L_Z] = rnd_user_hyperparams(Z);

    for i = 1:N
        jj = filterU{i}.j;
        kk = filterU{i}.k;
        ratings = filterU{i}.ratings - mean_rating;
        
        if ~isempty(ratings)
            QQ = Q(:,jj)*Q(:,jj)';
            L_Pi = L_P + alpha*QQ;
            L_Pi_ = inv(L_Pi);
            mu_Pi = L_Pi_*(L_P*mu_P + alpha*Q(:,jj)*(ratings - sum(W(:,jj).*X(:,kk),1)' - Y(:,kk)'*Z(:,i) ));
            P(:,i) = rnd_mvn(mu_Pi, L_Pi_);
        end
    end
    
    for j = 1:M
        ii = filterI{j}.i;
        kk = filterI{j}.k;
        ratings = filterI{j}.ratings - mean_rating;
        
        if ~isempty(ratings)
            PP = P(:,ii)*P(:,ii)';
            L_Qj = L_Q + alpha*PP;
            L_Qj_ = inv(L_Qj);
            mu_Qj = L_Qj_*(L_Q*mu_Q + alpha*P(:,ii)*(ratings - X(:,kk)'*W(:,j) - sum(Y(:,kk).*Z(:,ii),1)' ));
            Q(:,j) = rnd_mvn(mu_Qj, L_Qj_);
        end
    end
    
    
    for j = 1:M
        ii = filterI{j}.i;
        kk = filterI{j}.k;
        ratings = filterI{j}.ratings - mean_rating;
        
        if ~isempty(ratings)
            XX = X(:,kk)*X(:,kk)';
            L_Wj = L_W + alpha*XX;
            L_Wj_ = inv(L_Wj);
            mu_Wj = L_Wj_*(L_W*mu_W + alpha*X(:,kk)*(ratings - P(:,ii)'*Q(:,j) - sum(Y(:,kk).*Z(:,ii),1)' ));
            W(:,j) = rnd_mvn(mu_Wj, L_Wj_);
        end
    end
    
    for k = 1:K
        ii = filterT{k}.i;
        jj = filterT{k}.j;
        ratings = filterT{k}.ratings - mean_rating;
        
        if ~isempty(ratings)
            WW = W(:,jj)*W(:,jj)';
            WR = W(:,jj)*(ratings - sum(P(:,ii).*Q(:,jj),1)' - Z(:,ii)'*Y(:,k));
            switch k
                case 1,
                    L_Xk = 2*L_X + alpha*WW;
                    L_Xk_ = inv(L_Xk);
                    mu_Xk = (mu_X + X(:,2))*0.5 ;
                case K,
                    L_Xk = L_X + alpha*WW;
                    L_Xk_ = inv(L_Xk);
                    mu_Xk = L_Xk_*(L_X*X(:,K-1) + alpha*WR);
                otherwise
                    L_Xk = 2*L_X + alpha*WW;
                    L_Xk_ = inv(L_Xk);
                    mu_Xk = L_Xk_*(L_X*(X(:,k-1)+X(:,k+1)) + alpha*WR);
            end
            X(:,k) = rnd_mvn(mu_Xk, L_Xk_);
        end
    end
    
    for k = 1:K
        ii = filterT{k}.i;
        jj = filterT{k}.j;
        ratings = filterT{k}.ratings - mean_rating;
        
        if ~isempty(ratings)
            ZZ = Z(:,ii)*Z(:,ii)';
            ZR = Z(:,jj)*(ratings - sum(P(:,ii).*Q(:,jj),1)' - W(:,jj)'*X(:,k));
            switch k
                case 1,
                    L_Yk = 2*L_Y + alpha*ZZ;
                    L_Yk_ = inv(L_Yk);
                    mu_Yk = (mu_Y + Y(:,2))*0.5; 
                case K,
                    L_Yk = L_Y + alpha*ZZ;
                    L_Yk_ = inv(L_Yk);
                    mu_Yk = L_Yk_*(L_Y*Y(:,K-1) + alpha*ZR);
                otherwise
                    L_Yk = 2*L_Y + alpha*WW;
                    L_Yk_ = inv(L_Yk);
                    mu_Yk = L_Yk_*(L_Y*(Y(:,k-1)+Y(:,k+1)) + alpha*ZR);
            end
            Y(:,k) = rnd_mvn(mu_Yk, L_Yk_);
        end
    end
    
   for i = 1:N
        jj = filterU{i}.j;
        kk = filterU{i}.k;
        ratings = filterU{i}.ratings - mean_rating;
        
        if ~isempty(ratings)
            YY = Y(:,kk)*Y(:,kk)';
            L_Zi = L_Z + alpha*YY;
            L_Zi_ = inv(L_Zi);
            mu_Zi = L_Zi_*(L_Z*mu_Z + alpha*Y(:,kk)*(ratings - Q(:,jj)'*P(:,i) - sum(W(:,jj).*X(:,kk),1)' ));
            Z(:,i) = rnd_mvn(mu_Zi, L_Zi_);
        end
   end
   
	out_ratings = predict(test, P, Q, W, X, Y, Z);
	out_ratings = out_ratings + mean_rating; 
   	all_ratings = (all_ratings*count + out_ratings)/(count+1);   
   	rmse = sqrt(mean((all_ratings-double(test(:,4))).^2));   
	err(epoch) = rmse;		
	fprintf('iter - %3d, rmse = %2.4f \n',epoch, rmse);
	count = count + 1;
end
toc

BPTF2_Wt.P = P;
BPTF2_Wt.Q = Q;
BPTF2_Wt.W = W;
BPTF2_Wt.X = X;
BPTF2_Wt.Y = Y;
BPTF2_Wt.Z = Z;
BPTF2_Wt.err = err;

outfile = sprintf('%s_BPTF2_D%d.mat', tag, D);
save(outfile, 'BPTF2_Wt');	
